#include "mloop.h"
#include "grid.h"

#define HBOND1 2.70
#define HBOND2 3.00
#define SCALE14 0.0

#define PACKING_GAP 3
#define ENERGY_PENALTY 50.0

/* calc the energy of two atoms */
float calc_energy(float coor1[3], float coor2[3], unsigned char status1, unsigned char status2, float radius1, float radius2, struct atomname *atom1, struct atomname *atom2)
{
	float d, d2, energy;

	d = (coor1[0] - coor2[0]) * (coor1[0] - coor2[0]) + (coor1[1] - coor2[1]) * (coor1[1] - coor2[1]) + (coor1[2] - coor2[2]) * (coor1[2] - coor2[2]);
	d2 = (radius1 + radius2) * (radius1 + radius2);
	if (d >= d2)
		energy = 0.0;
	else
	{
		if (((status1 & IS_DONO) && (status2 & IS_ACCE)) || ((status1 & IS_ACCE) && (status2 & IS_DONO)))
		{
			if (((status1 & IS_DONO) && atom2->atomname[0] == 'C') || ((status2 & IS_DONO) && atom1->atomname[0] == 'C'))
				d2 = HBOND2 * HBOND2;
			else
				d2 = HBOND1 * HBOND1;
		}
		if (d < d2)
		{
			energy = K_OVERLAP * (d2 - d);
			//printf ("%s %s - %s %s %8.3f%8.3f%8.3f%8.3f%8.3f%8.3f%8.3f%8.3f\n", atom1->resname, atom1->atomname, atom2->resname, atom2->atomname, coor1[0], coor1[1], coor1[2], coor2[0], coor2[1], coor2[2], d, energy);
		}
		else
			energy = 0.0;
	}

	return energy;
}

/* calculate all energy */
int calc_energy1(struct candidate *p1)
{
	float x, y, z;
	int i, j, k, l, n;

	p1->energy1 = p1->energy2 = 0.0;
	p1->packing = 0;

	for (i = loop.start[0]; i <= loop.end[loop.resnum - 1]; i++)
	{
		if (loop.loopatom[i].build_flag == 0)
			continue;
		if (mode == 1 && !(loop.loopatom[i].status & IS_MAINCHAIN_AND_CB))
			continue;
		if (mode == 2 && !(loop.loopatom[i].status & IS_LOCAL_AA || loop.loopatom[i].status & IS_MAINCHAIN_AND_CB))
			continue;

		//printf ("%4d %s %s %8.3f%8.3f%8.3f\n", loop.loopatom[i].resno, loop.loopatom[i].name.resname, loop.loopatom[i].name.atomname, p1->coor[i][0], p1->coor[i][1], p1->coor[i][2]);
		x = p1->coor[i][0] - grid.center[0];
		y = p1->coor[i][1] - grid.center[1];
		z = p1->coor[i][2] - grid.center[2];

		x = GRIDMAX_X * 0.5 + x / GRID_SIZE;
		y = GRIDMAX_Y * 0.5 + y / GRID_SIZE;
		z = GRIDMAX_Z * 0.5 + z / GRID_SIZE;

		j = (int) x - (x > 0 ? 0 : 1);
		k = (int) y - (y > 0 ? 0 : 1);
		l = (int) z - (z > 0 ? 0 : 1);

		if (j < 0 || k < 0 || l < 0 || j >= GRIDMAX_X || k >= GRIDMAX_Y || l >= GRIDMAX_Z)
		{
			return -1;
		}

		//if (packing == 1 && (loop.loopatom[i].status & IS_CA) && grid2.cell2[j][k][l].flag == 0)
		if ((loop.loopatom[i].status & IS_CA) && grid2.cell2[j][k][l].flag == 0)
		{
			float dist;
        
			dist = 2 * PACKING_CUTOFF * PACKING_CUTOFF;
			for (n = 0; n < grid2.cell2[j][k][l].number; n++)
			{
				int *p2;
        
				p2 = grid2.cell2[j][k][l].p + n;
				if (environatom[*p2].resno - loop.loopatom[i].resno <= PACKING_GAP && environatom[*p2].resno - loop.loopatom[i].resno >= -PACKING_GAP)
					continue;
				dist = (p1->coor[i][0] - environatom[*p2].coor[0]) * (p1->coor[i][0] - environatom[*p2].coor[0]) + (p1->coor[i][1] - environatom[*p2].coor[1]) * (p1->coor[i][1] - environatom[*p2].coor[1]) + (p1->coor[i][2] - environatom[*p2].coor[2]) * (p1->coor[i][2] - environatom[*p2].coor[2]);
				if (dist < PACKING_CUTOFF * PACKING_CUTOFF)
					break;
			}

			if (dist > PACKING_CUTOFF * PACKING_CUTOFF)
			{
				for (n = 0; n < loop.loopnum; n++)
				{
					if (loop.loopatom[n].resno - loop.loopatom[i].resno <= PACKING_GAP && loop.loopatom[n].resno - loop.loopatom[i].resno >= -PACKING_GAP)
						continue;
					if (!(loop.loopatom[n].status & IS_MAINCHAIN_AND_CB))
						continue;
					dist = (p1->coor[i][0] - p1->coor[n][0]) * (p1->coor[i][0] - p1->coor[n][0]) + (p1->coor[i][1] - p1->coor[n][1]) * (p1->coor[i][1] - p1->coor[n][1]) + (p1->coor[i][2] - p1->coor[n][2]) * (p1->coor[i][2] - p1->coor[n][2]);
					if (dist < PACKING_CUTOFF * PACKING_CUTOFF)
						break;
				}
			}

			if (dist > PACKING_CUTOFF * PACKING_CUTOFF)
			{
				p1->energy1 += ENERGY_PENALTY * loop.resnum;
				p1->packing ++;
			}
		}

		for (n = 0; n < grid.cell[j][k][l].number; n++)
		{
			int *p2;

			p2 = grid.cell[j][k][l].p + n;
			p1->energy2 += loop.loopatom[i].scale_factor * calc_energy (p1->coor[i], environatom[*p2].coor, loop.loopatom[i].status, environatom[*p2].status, loop.loopatom[i].radius, environatom[*p2].radius, &(loop.loopatom[i].name), &(environatom[*p2].name));
		}

		for (j = 0; j < loop.loopnum; j++)
		{
			float tmp;

			if (loop.loopatom[j].build_flag == 1)
			{
				if (i >= j)
					continue;
				if (mode == 1 && !(loop.loopatom[j].status & IS_MAINCHAIN_AND_CB))
					continue;
				if (mode == 2 && !(loop.loopatom[j].status & IS_LOCAL_AA || loop.loopatom[j].status & IS_MAINCHAIN_AND_CB))
					continue;
			}

			if (table[i][j] == 0)
				tmp = 1.0;
			else if (table[i][j] == 1)
				continue;
			else if (table[i][j] == 2)
				tmp = SCALE14;

			p1->energy2 += tmp * loop.loopatom[i].scale_factor * loop.loopatom[j].scale_factor * calc_energy (p1->coor[i], p1->coor[j], loop.loopatom[i].status, loop.loopatom[j].status, loop.loopatom[i].radius, loop.loopatom[j].radius, &(loop.loopatom[i].name), &(loop.loopatom[j].name));
		}
	}

	p1->energy = p1->energy1 + p1->energy2;

	return 0;
}

/* calculate energy of sidechain */
float calc_energy2(struct candidate *p1, int resnum)
{
	float energy, x, y, z;
	int i, j, k, l, n;

	energy = 0.0;

	for (i = loop.start[resnum] + 5; i <= loop.end[resnum]; i++)
	{
		x = p1->coor[i][0] - grid.center[0];
		y = p1->coor[i][1] - grid.center[1];
		z = p1->coor[i][2] - grid.center[2];

		x = GRIDMAX_X * 0.5 + x / GRID_SIZE;
		y = GRIDMAX_Y * 0.5 + y / GRID_SIZE;
		z = GRIDMAX_Z * 0.5 + z / GRID_SIZE;

		j = (int) x - (x > 0 ? 0 : 1);
		k = (int) y - (y > 0 ? 0 : 1);
		l = (int) z - (z > 0 ? 0 : 1);

		if (j < 0 || k < 0 || l < 0 || j >= GRIDMAX_X || k >= GRIDMAX_Y || l >= GRIDMAX_Z)
		{
			return -1.0;
		}

		for (n = 0; n < grid.cell[j][k][l].number; n++)
		{
			int *p2;

			p2 = grid.cell[j][k][l].p + n;
			energy += loop.loopatom[i].scale_factor * calc_energy (p1->coor[i], environatom[*p2].coor, loop.loopatom[i].status, environatom[*p2].status, loop.loopatom[i].radius, environatom[*p2].radius, &(loop.loopatom[i].name), &(environatom[*p2].name));
		}

		for (j = 0; j < loop.loopnum; j++)
		{
			float tmp;

			if (loop.loopatom[j].build_flag == 1)
			{
				if (loop.loopatom[j].resno == loop.loopatom[i].resno && i >= j)
					continue;
				if (mode == 1 && !(loop.loopatom[j].status & IS_MAINCHAIN_AND_CB))
					continue;
				if (mode == 2 && !(loop.loopatom[j].status & IS_LOCAL_AA || loop.loopatom[j].status & IS_MAINCHAIN_AND_CB))
					continue;
			}

			if (table[i][j] == 0)
				tmp = 1.0;
			else if (table[i][j] == 1)
				continue;
			else if (table[i][j] == 2)
				tmp = SCALE14;

			energy += tmp * loop.loopatom[i].scale_factor * loop.loopatom[j].scale_factor * calc_energy (p1->coor[i], p1->coor[j], loop.loopatom[i].status, loop.loopatom[j].status, loop.loopatom[i].radius, loop.loopatom[j].radius, &(loop.loopatom[i].name), &(loop.loopatom[j].name));
		}
	}

	return energy;
}
